from pathlib import Path

from datasets import load_dataset,ClassLabel
from transformers import AutoTokenizer
from configuration import config

def process_data():
    print("开始处理数据")
    #读取文件
    dataset_dict = load_dataset('csv',data_files={
        'train':str(config.RAW_DATA_DIR / 'train.txt'),
        'test': str(config.RAW_DATA_DIR / 'test.txt'),
        'valid': str(config.RAW_DATA_DIR / 'valid.txt'),
    },delimiter='\t')

    #过滤数据
    dataset_dict=dataset_dict.filter(lambda x:x['label'] is not None and x['text_a'] is not None)

    model_path = Path("D:/product-classify-bert/models/bert-base-chinese")
    #构建数据集
    tokenizer = AutoTokenizer.from_pretrained(str(model_path))

    #统计标题的长度
    df = dataset_dict['train'].to_pandas()
    print(df['text_a'].apply(lambda x: len(tokenizer.tokenize(x))).max())

    def tokenize(batch):
        tokenized = tokenizer(batch['text_a'],padding='max_length',truncation=True,max_length=config.SEQ_LEN)
        return {'input_ids': tokenized['input_ids'],'attention_mask': tokenized['attention_mask']}
    """
    该函数用于对一批文本输入进行分词，以便进行自然语言处理任务。

    此函数处理一批文本输入，将它们转换为适合NLP模型的格式。
    它使用预定义的分词器将文本转换为标记（tokens），
    并将这些标记填充（pad）或截断（truncate）到指定的长度。
    此外，它还会生成注意力掩码（attention masks），以指示哪些标记对模型的注意力机制是重要的。

    参数：
    - batch: 一个字典，包含需要分词的文本，其中键 'text_a' 指向文本输入的列表。

    返回：
    - 一个字典，包含两个键：
        * 'input_ids': 分词后并经过填充或截断的输入ID。
        * 'attention_mask': 与输入ID对应的注意力掩码。
    """
    # 对数据集中的每个样本应用 tokenize 函数进行文本编码处理，启用批量处理模式，并移除原始文本列 'text_a'
    dataset_dict = dataset_dict.map(tokenize, batched=True, remove_columns=['text_a'])

    # 处理标签列，获取训练集中所有唯一的标签类别
    all_labels = dataset_dict['train'].unique('label')
    class_label = ClassLabel(names=all_labels)
    # 打印类别总数
    print('类别总类：', len(all_labels))
    # 将数据集中的 'label' 列转换为 `ClassLabel` 类型，便于后续模型处理
    dataset_dict = dataset_dict.cast_column("label", ClassLabel(num_classes=len(all_labels), names=all_labels))

    #保存数据集
    dataset_dict['train'].save_to_disk(config.PROCESSED_DATA_DIR / 'train')
    dataset_dict['test'].save_to_disk(config.PROCESSED_DATA_DIR / 'test')
    dataset_dict['valid'].save_to_disk(config.PROCESSED_DATA_DIR / 'valid')

    print("处理数据完成")